import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import RegularPolygon
from matplotlib.path import Path
from matplotlib.projections.polar import PolarAxes
from matplotlib.projections import register_projection
from matplotlib.spines import Spine
from matplotlib.transforms import Affine2D


# 定义雷达图投影
def radar_factory(num_vars, frame='polygon'):
    theta = np.linspace(0, 2 * np.pi, num_vars, endpoint=False)

    class RadarAxes(PolarAxes):
        name = 'radar'

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self.set_theta_zero_location('N')

        def fill(self, *args, closed=True, **kwargs):
            return super().fill(closed=closed, *args, **kwargs)

        def plot(self, *args, **kwargs):
            lines = super().plot(*args, **kwargs)
            for line in lines:
                self._close_line(line)

        def _close_line(self, line):
            x, y = line.get_data()
            if x[0] != x[-1]:
                x = np.concatenate((x, [x[0]]))
                y = np.concatenate((y, [y[0]]))
                line.set_data(x, y)

        def set_varlabels(self, labels):
            self.set_thetagrids(np.degrees(theta), labels)

        def _gen_axes_patch(self):
            if frame == 'polygon':
                return RegularPolygon((0.5, 0.5), num_vars,
                                      radius=.55, edgecolor="none")
            else:
                raise ValueError("unknown value for 'frame': %s" % frame)

        def _gen_axes_spines(self):
            if frame == 'polygon':
                spine = Spine(axes=self,
                              spine_type='circle',
                              path=Path.unit_regular_polygon(num_vars))
                spine.set_transform(Affine2D().scale(.55).translate(.5, .5)
                                    + self.transAxes)
                spine.set_color('none')
                return {'polar': spine}
            else:
                raise ValueError("unknown value for 'frame': %s" % frame)

    register_projection(RadarAxes)
    return theta


# 调整categories顺序，让'Age\nImplicit'放在第一位（最上方）
# 原顺序: [0:EmotionExplicit, 1:EmotionImplicit, 2:AgeExplicit, 3:AgeImplicit, ...]
# 顺时针转3位后新顺序: [3:AgeImplicit, 4:GenderExplicit, 5:GenderImplicit, 6:EventExplicit, 7:EventImplicit, 0:EmotionExplicit, 1:EmotionImplicit, 2:AgeExplicit]
categories = [
    'Age\nImplicit',
    'Gender\nExplicit', 'Gender\nImplicit',
    'Event\nExplicit', 'Event\nImplicit',
    'Emotion\nExplicit', 'Emotion\nImplicit',
    'Age\nExplicit'
]
N = len(categories)
theta = radar_factory(N, frame='polygon')

# 相应调整每个模型的数据顺序，保持与categories对应
models = {
    'qwen2.5-omni': [55, 0, 27.80, 37, 3, 39.50, 42.41, 1],
    'freeze-omni': [48.40, 4, 25.60, 1, 2, 28.20, 41.8, 2],
    'GLM-4-voice': [52, 0, 25, 0, 0, 58.50, 47.4, 3],
    'baichuan-audio': [43.20, 0, 26, 7, 0, 38.80, 39.6, 1],
    'kimi-audio': [45.40, 62, 23.20, 12, 2, 58.50, 44, 39],
    'opens2s': [69.6, 46.5, 27.6, 6, 7, 41.5, 46.6, 30],
    'OSUM-chat': [75.20, 89, 57.40, 53.50, 95, 69.60, 58.3, 72.5]
}

# 计算每个维度的实际最大值
data_values = np.array(list(models.values()))
max_per_dimension = np.max(data_values, axis=0)

# 标准化数据
normalized_models = {}
for model, data in models.items():
    normalized = [d / max_val if max_val != 0 else 0 for d, max_val in zip(data, max_per_dimension)]
    normalized_models[model] = normalized

# 绘图设置
fig, ax = plt.subplots(figsize=(15, 15), subplot_kw=dict(projection='radar'))
ax.set_ylim(0, 1.1)

# 颜色方案
colors = [
    '#4285F4', '#34A853', '#FBBC05',
    '#EA4335', '#9C27B0', '#FF9800', '#00ACC1'
]

# 绘制雷达图并标注数值
for (model, normalized_data), (original_data, color) in zip(normalized_models.items(), zip(models.values(), colors)):
    ax.plot(theta, normalized_data, color=color, linewidth=2.5, label=model)
    ax.fill(theta, normalized_data, facecolor=color, alpha=0.05)

    # 标注原始数值
    for angle, norm_val, orig_val in zip(theta, normalized_data, original_data):
        x = angle
        y = norm_val + 0.05
        ax.text(x, y, f'{orig_val:.1f}',
                ha='center', va='center',
                color=color, fontweight='bold',
                fontsize=10)

# 设置坐标轴标签
ax.set_varlabels(categories)
labels = ax.get_xticklabels()
for label in labels:
    label.set_fontsize(12)
    label.set_fontweight('bold')
    label.set_color('#666666')

# 清除默认刻度
ax.set_rticks([])

# 绘制从中心到顶点的轴线
for angle in theta:
    ax.plot([angle, angle], [0, 1.2], color='gray', linestyle='-', linewidth=0.8)

# 添加网格线（环形）
ring_values = [0.3, 0.6, 0.9]
for ring in ring_values:
    ax.plot(theta, [ring] * N, color='gray', linestyle='--', linewidth=0.5)

# 调整图例位置
ax.legend(loc='upper right', bbox_to_anchor=(1.3, 1.1), fontsize=12)

plt.title("", pad=20, fontsize=16, fontweight='bold')
plt.tight_layout()
plt.show()
